package com.xfcode.ai.platform;

import com.google.gson.Gson;
import com.xfcode.ai.platform.bean.*;
import io.reactivex.BackpressureStrategy;
import io.reactivex.Flowable;
import okhttp3.*;
import org.jetbrains.annotations.NotNull;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

public class DeepSeekClient {
    private static final String DEFAULT_BASE_URL = "https://api.deepseek.com";
    private static final String CHAT_COMPLETION_SUFFIX = "/chat/completions";

    private String apiKey;

    public DeepSeekClient(String apiKey) {
        this.apiKey = apiKey;
    }
    public String getChatCompletionUrl() {
        return DEFAULT_BASE_URL + CHAT_COMPLETION_SUFFIX;
    }
    public String getApiKey() {
        return apiKey;
    }

    public String call(String template, String input) {
        List<ChatCompletionMessage> messages=new ArrayList<>();
        messages.add(new ChatCompletionMessage(ChatMessageRole.SYSTEM.value(),template));
        messages.add(new ChatCompletionMessage(ChatMessageRole.USER.value(), input));
        ChatCompletionRequest chatCompletionRequest = getChatCompletionRequest(messages, 20000);
        try {
            return chatCompletion(chatCompletionRequest).getChoices().get(0).getMessage().getContent();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public Flowable<ChatCompletionStreamResponse> stream(String message) {
        List<ChatCompletionMessage> messages = new ArrayList<>();
        messages.add(new ChatCompletionMessage(ChatMessageRole.USER.value(), message));
        ChatCompletionRequest chatCompletionRequest =getChatCompletionRequest(messages, 2000);
        try {
            return  chatCompletionStream(chatCompletionRequest);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
    public Flowable<ChatCompletionStreamResponse> chatCompletionStream(ChatCompletionRequest request) throws IOException {
        request.stream = true;
        String content= new Gson().toJson(request);
        return getChatCompletionStreamResponseFlowable(content);
    }

    private @NotNull Flowable<ChatCompletionStreamResponse> getChatCompletionStreamResponseFlowable( String content) throws IOException {
        OkHttpClient client = new OkHttpClient();
        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType,content);
        Request httpRequest = new Request.Builder()
                .url(getChatCompletionUrl())
                .addHeader("Authorization", "Bearer " + getApiKey() )
                .addHeader("Content-Type", "application/json")
                .post(body)
                .build();
        Response response = client.newCall(httpRequest).execute();
        if (response.code() != 200) {
            throw new RuntimeException("Failed to start stream: " + response.body().string());
        }

        // get response body line by line
        return Flowable.create(emitter -> {
            ResponseBody responseBody = response.body();
            if (responseBody == null) {
                emitter.onError(new RuntimeException("Response body is null"));
                return;
            }
            String line;
            while ((line = responseBody.source().readUtf8Line()) != null) {
                if (line.startsWith("data:")) {
                    line = line.substring(5);
                    line = line.trim();
                }
                if (Objects.equals(line, "[DONE]")) {
                    emitter.onComplete();
                    return;
                }
                line = line.trim();
                if (line.isEmpty()) {
                    continue;
                }
                Gson gson = new Gson();
                ChatCompletionStreamResponse streamResponse = gson.fromJson(line, ChatCompletionStreamResponse.class);
                emitter.onNext(streamResponse);
            }
            emitter.onComplete();
        }, BackpressureStrategy.BUFFER);
    }

    private static @NotNull ChatCompletionRequest getChatCompletionRequest(List<ChatCompletionMessage> messages, int maxTokens) {
        return new ChatCompletionRequest(
                "deepseek-chat",
                messages,
                maxTokens,
                1.3f,
                1
        );
    }

    public ChatCompletionResponse chatCompletion(ChatCompletionRequest request) throws IOException {
        request.stream = false;
        okhttp3.OkHttpClient client = new okhttp3.OkHttpClient();
        okhttp3.MediaType mediaType = okhttp3.MediaType.parse("application/json");
        okhttp3.RequestBody body = okhttp3.RequestBody.create(mediaType, new Gson().toJson(request));
        okhttp3.Request httpRequest = new okhttp3.Request.Builder()
                .url(getChatCompletionUrl())
                .addHeader("Authorization", "Bearer " + getApiKey() )
                .addHeader("Content-Type", "application/json")
                .post(body)
                .build();
        try {
            okhttp3.Response response = client.newCall(httpRequest).execute();
            String responseBody = response.body().string();
            Gson gson = new Gson();
            return gson.fromJson(responseBody, ChatCompletionResponse.class);
        } catch (IOException e) {
            e.printStackTrace();
            throw e;
        }
    }

}
